import os
import random
import torch
import pandas as pd
import numpy as np
from model import TransformerAutoEncoder
from torch.utils.data import TensorDataset, DataLoader
from sklearn.model_selection import train_test_split
from sklearn.preprocessing import MinMaxScaler
from sklearn.metrics import recall_score
from sklearn.linear_model import LogisticRegression
from sklearn.tree import DecisionTreeClassifier
from sklearn.ensemble import RandomForestClassifier
import joblib


# Set env
def seed_torch(seed=42):
    random.seed(seed)
    os.environ['PYTHONHASHSEED'] = str(seed)
    np.random.seed(seed)
    torch.manual_seed(seed)
    torch.cuda.manual_seed(seed)
    torch.cuda.manual_seed_all(seed)  # if you are using multi-GPU.
    torch.backends.cudnn.benchmark = False
    torch.backends.cudnn.deterministic = True
    torch.backends.cudnn.enabled = True


seed_torch()

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

# Load data
train_df = pd.read_csv('dataset/cs-training.csv')
median = train_df.median()
train_df = train_df.fillna(median)
test_df = pd.read_csv('dataset/cs-test.csv')
test_df = test_df.fillna(median)

print(train_df.columns)
train = train_df.values[:, 1:]
test = test_df.values[:, 1:]

X, y = train[:, 1:], train[:, 0]
X_test, y_test = test[:, 1:], test[:, 0]

# Min-max scaler
scaler = MinMaxScaler().fit(X)
X = scaler.transform(X)
X_test = scaler.transform(X_test)

X_train, X_val, y_train, y_val = train_test_split(X, y, test_size=0.2, random_state=42, stratify=y)

# Load features
train_features = joblib.load('train_features.npy')
train_features_c = joblib.load('train_features_c.npy')
val_features = joblib.load('val_features.npy')
test_features = joblib.load('test_features.npy')

# LogisticRegression
lr = LogisticRegression(random_state=42, class_weight={0: 1, 1: 7}, max_iter=300, tol=1e-3).fit(X_train, y_train)
lr_preds = lr.predict(X_val)
lr_recall_o = recall_score(y_val, lr_preds)
print('lr_recall:', lr_recall_o)

lr = LogisticRegression(random_state=42, class_weight={0: 1, 1: 3}, max_iter=300, tol=1e-3, solver='liblinear').fit(
    train_features, y_train)
lr_preds = lr.predict(val_features)
lr_recall = recall_score(y_val, lr_preds)
print('transformer_lr_recall:', lr_recall)

lr = LogisticRegression(random_state=42, class_weight={0: 1, 1: 3}, max_iter=300, tol=1e-3, solver='liblinear').fit(
    train_features_c, y_train)
lr_preds = lr.predict(val_features)
lr_recall_c = recall_score(y_val, lr_preds)
print('transformer_lr_recall_c:', lr_recall_c)

with open(f"log/{os.path.basename(__file__)}.log", "a") as f:
    f.write(f"lf_recall: {lr_recall_o}, "
            f"transformer_lf_recall: {lr_recall}, "
            f"transformer_lf_recall_c: {lr_recall_c}"
            f"\n")